import torch
import pandas as pd
import math
from src.utils.log import infomsg, errormsg


def load_space_data(path, num_nodes, mode, para_lambda=500, limit=2000):
    """
    Args:
        path (str): 空间文件的路径
        num_nodes (int): 风机的数量
        para_lambda (float): lambda参数
        limit (float): 高于limit的距离将会被过滤掉
    
    Returns: 
        torch.Tensor: 空间相似度矩阵，尺寸为(N, N)
    """
    data = None
    if path is not None:
        data = pd.read_csv(path)
    space_graph = generate_adj(data, num_nodes, para_lambda, limit, mode).float()
    return space_graph


def generate_adj(data, num_nodes, para_lambda, limit, mode='lal'):
    '''
    根据每台风机的位置信息，调用公式为每对风机生成(0, 1)之间的相似度
    公式为 w = e ^ (- (d / λ) ^ 2)，其中d是风机之间的相对距离
    
    Args:
        data (pd.DataFrame): 多台风机的位置信息
        para_lambda (float): λ参数
        limit (float): 限制值
        mode (str): 距离信息的给出方式，lal表示'longitude and latitude'，即经纬度，'coor'表示直角坐标，'None'表示没有距离信息
    
    Returns: 
        torch.Tensor: 相似度矩阵
    '''
    
    if mode == 'None':
        return torch.eye(num_nodes, dtype=torch.float32)
    # 生成矩阵
    matrix = pd.DataFrame(index=data.index, columns=data.index)
    for i in data.index:
        for j in data.index:
            if i == j:
                matrix.loc[i, j] = 0
            else:
                if mode == 'lal':
                    matrix.loc[i, j] = distance_lal(data.loc[i, 'longitude'], data.loc[i, 'latitude'], data.loc[j, 'longitude'], data.loc[j, 'latitude'])
                elif mode == 'coor':
                    matrix.loc[i, j] = distance_coor(data.loc[i, 'x'], data.loc[i, 'y'], data.loc[j, 'x'], data.loc[j, 'y'])
                dis = matrix.loc[i, j]
                if dis < limit:
                    matrix.loc[i, j] = math.exp(-dis ** 2 / para_lambda ** 2)
                else:
                    matrix.loc[i, j] = 0
            
    matrix = matrix.astype(float)
    matrix = torch.tensor(matrix.values, dtype=torch.float32)
    return matrix


def distance_coor(x1, y1, x2, y2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5


def distance_lal(lon1: float, lat1: float, lon2: float, lat2: float) -> float:
    '''
    给出两对经纬度，计算以米为单位的距离

    Args: 
        lon1 (float): 经度1
        lat1 (float): 纬度1
        lon2 (float): 经度2
        lat2 (float): 纬度2
    
    Returns:
        float: 距离
    '''

    # 地球半径
    R = 6371000  
    
    # 将经纬度转换为弧度制
    lat1_rad = math.radians(lat1)
    lon1_rad = math.radians(lon1)
    lat2_rad = math.radians(lat2)
    lon2_rad = math.radians(lon2)
    
    # 计算经纬度的差值
    dlat = lat2_rad - lat1_rad
    dlon = lon2_rad - lon1_rad
    
    # 使用 Haversine 公式计算距离
    a = math.sin(dlat / 2) ** 2 + math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(dlon / 2) ** 2
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    ans = R * c
    return ans